import pandas as pd
import os
import torch
from torch.utils.data import Dataset,DataLoader
import glob
from PIL import Image
import torchvision.transforms.functional as ff
import numpy as np
import torchvision.transforms as transforms


class CamvidDataset(Dataset):
    def __init__(self,file_path = [],crop_size=(352, 480),is_train = True,transform =None):
        super(CamvidDataset, self).__init__()
        self.is_train = is_train
        self.image_path,self.label_path = file_path
        self.images,self.labels = self.read_image_label(self.image_path,self.label_path)
        self.transform = transform
        self.crop_size = crop_size

    def __len__(self):
        return len(self.images)
    def __getitem__(self, index):
        image_path = self.images[index]
        image = Image.open(image_path)
        image = ff.center_crop(image,self.crop_size)
        if self.is_train:
            label_path = self.labels[index]
            label = Image.open(label_path).convert('RGB')
            label = ff.center_crop(label,self.crop_size)
            encode_label = label_processor.encode_label_img(label)
            encode_label = torch.from_numpy(encode_label.copy())
            if self.transform:
                image = self.transform(image)
            sample = {'image':image,'mask':encode_label}
            return sample

        else:
            return self.transform(image)
    def read_image_label(self,image_path,label_path):
        image_names = os.listdir(image_path)
        images = [os.path.join(image_path, image_name) for image_name in image_names]
        # print(self.is_train)
        if self.is_train:
            label_names = [(path.split('.')[0]+ "_L.png") for path in image_names]
            labels = [os.path.join(label_path, label_name) for label_name in label_names]
            return images,labels
        else:
            return images,None
class LabelProcessor:

    def __init__(self,file_path):
        self.colormap = self.read_color_map(file_path)
        self.cm2lbl = self.encode_label_pix(self.colormap)

    @staticmethod
    def read_color_map(file_path):
        label_color = pd.read_csv(file_path,sep =',')
        color_map = []
        for i in range(len(label_color.index)):
            temp = label_color.iloc[i]
            color = [temp['r'],temp['g'],temp['b']]
            color_map.append(color)
        return color_map
    @staticmethod
    def encode_label_pix(color_map):
        cm2lbl = np.zeros(256**3)
        for i,color in enumerate(color_map):
            cm2lbl[(color[0]*256+color[1])*256+color[2]] = i
        return cm2lbl
    def encode_label_img(self,img):
        data = np.array(img,dtype='int32')
        idx = (data[...,0]*256+data[...,1])*256+data[...,2]
        return np.array(self.cm2lbl[idx], dtype='int64')

label_processor =LabelProcessor('./CamVid/class_dict.csv')
# if __name__ == '__main__':
#     train_path = './CamVid/train'
#     label_path = './CamVid/train_labels'
#     transfrom = transforms.Compose(
#         [
#             transforms.ToTensor(),
#             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
#         ]
#     )
#     Cam_train_dataset = CamvidDataset([train_path,label_path],is_train=True,transform=transfrom)
#     dataloder = DataLoader(Cam_train_dataset,batch_size=1)
#     for img,label in dataloder:
#         print(label[0][:5,0:5])
#         print(img.shape,label.shape)
#         break







